package com.kingsoft.dc.khaos.scheduler

import java.util.concurrent.atomic.AtomicInteger

import com.kingsoft.dc.khaos.util.{KhaosConstants, Logging}
import com.kingsoft.dc.khaos.metadata.{ModuleCommon, TaskInfo}
import com.kingsoft.dc.khaos.KhaosContext
import org.apache.spark.util.LongAccumulator

import scala.collection.immutable.HashMap
import scala.collection.mutable
import scala.collection.mutable.ListBuffer

/**
  * Created by chengguangqing on 2019/5/27.
  */
private[khaos] class KhaosScheduler(private[scheduler] val kc:KhaosContext,
                             private[scheduler] val jobScheduler:JobScheduler,
                             private[scheduler] val taskScheduler: TaskScheduler
                               //listenerBus:LiveListenerBus 21202F2938212B3E22272626252E434D
                             ) extends Logging{

  def this(kc:KhaosContext){
    this(kc,kc.jobScheduler,kc.taskScheduler)
  }

  private[scheduler] val nextApplicationId = new AtomicInteger(0)
  private[scheduler] def numTotalApplications: Int = nextApplicationId.get()
  private[scheduler] val applicationIdToJobs = new HashMap[Int,KhaosJob]()

  private[scheduler] val pendingTasks = new mutable.HashSet[Int]()
  private[scheduler] val failedTasks = new mutable.HashSet[Int]()
  private[scheduler] val runningTasks = new mutable.HashSet[Int]()
  private[scheduler] val completedTasks = new mutable.HashSet[Int]()
  private[scheduler] val activeTasks = new mutable.HashSet[Int]()

  private[scheduler] val activeJobs = new mutable.HashSet[Int]
  private[scheduler] val failedJobs = new mutable.HashSet[Int]()
  private[scheduler] val completedJobs = new mutable.HashSet[Int]()

  private[scheduler] val activeApplications = new mutable.HashSet[Int]()
  private[scheduler] val failedApplications = new mutable.HashSet[Int]()
  private[scheduler] val completedApplications = new mutable.HashSet[Int]()

  private[scheduler] val eventProcessLoop = new KhaosSchedulerEventProcessLoop(this)

  private[scheduler] var mcDegreeOut: scala.collection.mutable.Map[String, ListBuffer[String]] = _

  jobScheduler.setKhaosScheduler(this)
  taskScheduler.setKhaosScheduler(this)

  private var jobNum: Int = 1

  //运行
  def run(modules: List[ModuleCommon]):Unit={

    val start = System.nanoTime()
    val waiter = applicationSubmitted(modules)

    waiter.awaitResult() match{
      case ApplicationSucceeded =>

        log.info("application %d finished: %s, took %f s".format
        (waiter.applicationId,"abcde",(System.nanoTime()-start/1e9)))

      case ApplicationFailed(exception:Exception) =>

        log.info("application %d finished: %s, took %f s".format
        (waiter.applicationId,"",(System.nanoTime()-start/1e9)))
        val callerStackTrace = Thread.currentThread().getStackTrace.tail
        exception.setStackTrace(exception.getStackTrace ++ callerStackTrace)
        throw exception
    }
  }

  //提交app
  def applicationSubmitted(modules: List[ModuleCommon]): ApplicationWaiter ={
    val applicationId = nextApplicationId.getAndIncrement()
    activeApplications += applicationId
    val waiter = new ApplicationWaiter(this, applicationId,jobNum)
    eventProcessLoop.post(ApplicationSubmitted(applicationId,modules,waiter))
    waiter
  }

  //提交Job
  def jobSubmitted(jobId: Int,
                   modules: List[ModuleCommon],
                   degreeOut: collection.mutable.Map[String, ListBuffer[String]],
                   applicationListener: ApplicationListener): Unit = {
    mcDegreeOut = degreeOut
    eventProcessLoop.post(JobSubmitted(jobId, modules, applicationListener))
  }

  //提交task
  def taskSubmitted(taskId: Int,
                    moduleCommon: ModuleCommon,
                    taskInfo: TaskInfo,
                    listener: ApplicationListener): Unit = {
    eventProcessLoop.post(TaskSubmitted(taskId, moduleCommon, taskInfo, listener))
  }

  def taskResultCollector(taskInfo: TaskInfo,
                          listener:ApplicationListener): Unit ={
    eventProcessLoop.post(TaskResultCollector(taskInfo,listener))
  }




  //handle
  def handleApplicationSubmitted(applicationId:Int,
                                 modules:List[ModuleCommon],
                                 listener:ApplicationListener):Unit={
    val applicationSubmissionTime = System.currentTimeMillis()

    log.info("调用handleApplicationSubmitted:"+applicationId)

    //添加到监听里面
    //listenerBus.post()
    ///添加 job 拆分业务逻辑

    try{
      //println(jobScheduler)
      this.jobScheduler.runJob(applicationId,modules,listener);
    } catch {
      case e:Exception =>
        //添加到监听中
        log.warn("Khaos Exception :" + e)
    }
  }

  def handleApplicationCompleted(applicationId: Int,
                                 completedTime: Long,
                                 applicationResult: ApplicationResult,
                                 listener: ApplicationListener): Unit = {
    //listenerBus.post()
    applicationResult match {
      case ApplicationSucceeded => listener.applicationSucceeded()
      case ApplicationFailed(ex) => listener.applicationFailed(ex)
    }
  }

  def cancelApplication(applicationId: Int): Unit = {
    log.info("Asked to cancel job " + applicationId)
    eventProcessLoop.post(ApplicationCancelled(applicationId))
  }


  def handleJobSubmitted(jobId: Int,
                         modules: List[ModuleCommon],
                         listener: ApplicationListener): Unit = {
    val jobSubmissionTime = System.currentTimeMillis()
     activeJobs += jobId
    //listenerBus.post()
    this.taskScheduler.submitTasks(jobId, modules, listener)
  }

  def handleJobCompleted(jobId: Int,
                         completedTime: Long,
                         jobResult: JobResult,
                         listener: ApplicationListener): Unit = {
    //listenerBus.post(KhaosListenerJobEnd(jobId, completedTime, jobResult, listener))
  }


  def handleTaskSubmitted(taskId: Int,
                          moduleCommon: ModuleCommon,
                          taskInfo: TaskInfo,
                          listener: ApplicationListener
                         ): Unit = {

    val taskSubmissionTime = System.currentTimeMillis()
    taskInfo.submissionTime = Some(taskSubmissionTime)

    runningTasks += taskId
    //listenerBus.post()
    if(moduleCommon.engine.equals("SPARK")){
      this.taskScheduler.runTask(taskId, moduleCommon, taskInfo, listener)
    }
    else{
      //其他引擎
    }
  }

  def cancelJob(jobId: Int): Unit = {
    log.info("Asked to cancel job " + jobId)
    eventProcessLoop.post(JobCancelled(jobId))
  }

  def handleTaskResultCollector(taskInfo: TaskInfo,
                                listener: ApplicationListener): Unit ={
    //listenerBus.post()
    mcDegreeOut.foreach(h => mcDegreeOut.get(h._1).get.-=(taskInfo.moduleId.toString))
    completedTasks += taskInfo.taskId
    runningTasks -= taskInfo.taskId
    pendingTasks -= taskInfo.taskId
    this.isLastTask(taskInfo,listener,null);
  }

  //任务失败
  def handleTaskFailed(taskInfo: TaskInfo,
                       exception: Option[Exception]): Unit = {
    failedTasks += taskInfo.taskId
    logError("task %s failed, message: %s".format(taskInfo.taskId, taskInfo.failureReason))

    //throw exception.get
  }


  def taskFailed(taskInfo: TaskInfo,
                 exception: Option[Exception],
                 jobId: Int,
                 listener: ApplicationListener): Unit ={
    //任务失败清理
    eventProcessLoop.onClear()
    pendingTasks.clear()
    eventProcessLoop.post(TaskFailed(taskInfo, exception))
    applicationCompletedUpdate(1, taskInfo.completionTime, ApplicationFailed(exception.getOrElse(new Exception("Khaos ApplicationFailed......"))), listener)

    /*failedTasks += taskInfo.taskId
    runningTasks -= taskInfo.taskId
    isLastTask(taskInfo,listener,exception)*/

  }

  //运行状态更新
  def isLastTask(taskInfo: TaskInfo,
                 listener: ApplicationListener,
                 exception: Option[Exception]): Unit = {

    //检查是否最有一个task
    if (pendingTasks.isEmpty) {
      val jobId = taskInfo.jobId
      //更新状态job状态
      completedJobs += jobId
      activeJobs -= jobId
      jobCompletedUpdate(jobId, taskInfo.completionTime, JobSucceeded, listener)

      if (activeJobs.isEmpty) {
        if(failedTasks.isEmpty){
          completedApplications += 1
          activeApplications -= 1
          applicationCompletedUpdate(1, taskInfo.completionTime, ApplicationSucceeded, listener)
        }
        else{
          completedApplications += 1
          activeApplications -= 1
          applicationCompletedUpdate(1, taskInfo.completionTime, ApplicationFailed(exception.getOrElse(new Exception("Khaos ApplicationFailed......"))), listener)
        }
      }
    }
  }


  //更新Job完成状态
  def jobCompletedUpdate(jobId: Int,
                   completedTime: Option[Long],
                   jobResult: JobResult,
                   listener: ApplicationListener): Unit = {
    eventProcessLoop.post(
      JobCompleted(jobId, completedTime.getOrElse(0L), jobResult, listener)
    )
  }

  //更新app完成状态
  def applicationCompletedUpdate(applicationId: Int,
                           completedTime: Option[Long],
                           applicationResult: ApplicationResult,
                           listener: ApplicationListener): Unit = {
    eventProcessLoop.post(
      ApplicationCompleted(applicationId, completedTime.getOrElse(0L), applicationResult, listener)
    )
  }

  //启动处理线程
  eventProcessLoop.start()
}
